
#include "../aprtools/APRBucketBrigade.h"
#include "testEncDecWrapper.h"
#include "../crypt/DSAKeyPairGenerator.h"
#include "../crypt/RSAKeyPairGenerator.h"

void tests::testEncDecWrapper(apr_pool_t *mp)
{
   testEncDecWrapperBlock(10, 10000, mp);
   testEncDecWrapperPipe(10, 10000, mp);
}

void tests::testEncDecWrapperPipe(int iterations, int max_length, apr_pool_t *mp)
{
   DSAKeyPairGenerator dsakpgen(DSA_BITS, "DSAkey");
   DSAPublicKey *d_pubkey = dsakpgen.popPublicKey();
   DSAPrivateKey *d_prvkey = dsakpgen.popPrivateKey();
   RSAKeyPairGenerator rsakpgen(RSA_BITS, "RSAkey");
   RSAPublicKey *r_pubkey = rsakpgen.popPublicKey();
   RSAPrivateKey *r_prvkey = rsakpgen.popPrivateKey();
   for (int i = 0; i < iterations; i++)
   {
      APRBucketBrigade *t_input = new APRBucketBrigade(mp);
      
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
      t_input->writeToStream(pseudorandom, randlen);
      t_input->writeEOS();
   
      EncryptWrapper *encwrap = new EncryptWrapper((EncryptionKey *)r_pubkey, (SigningKey *)d_prvkey, 1);
      int testtype = rand() % 2;
      DecryptWrapper *decwrap = new DecryptWrapper((DecryptionKey *)r_prvkey, (testtype) ? ((CheckingKey *)d_pubkey) : ((CheckingKey *)r_pubkey), 1);

      APRBucketBrigade *t_intermediate = encwrap->encryptToBucket(t_input);
      
      APRBucketBrigade *t_output = decwrap->decryptToBucket(t_intermediate);
      
      size_t new_plain_size = 0;
      char *new_plain_data = NULL;
      t_output->bucketBrigadeToBuffer(&new_plain_data, &new_plain_size);
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
      } else
      {
         printf("EncDecWrapper Pipe(issame) test %d passed - data length %d\n", i, randlen);
      }
      free(pseudorandom);
      free(new_plain_data);
      
      // Can't delete memory pools before deleting things in the memory pools.
      delete t_output;
      delete t_intermediate;
      delete decwrap;
      delete encwrap;
      delete t_input;
   }
   delete d_pubkey; delete d_prvkey;
   delete r_pubkey; delete r_prvkey;
}

void tests::testEncDecWrapperBlock(int iterations, int max_length, apr_pool_t *mp)
{
   DSAKeyPairGenerator dsakpgen(DSA_BITS, "DSAkey");
   DSAPublicKey *d_pubkey = dsakpgen.popPublicKey();
   DSAPrivateKey *d_prvkey = dsakpgen.popPrivateKey();
   RSAKeyPairGenerator rsakpgen(RSA_BITS, "RSAkey");
   RSAPublicKey *r_pubkey = rsakpgen.popPublicKey();
   RSAPrivateKey *r_prvkey = rsakpgen.popPrivateKey();
   for (int i = 0; i < iterations; i++)
   {
      
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
   
      EncryptWrapper *encwrap = new EncryptWrapper((EncryptionKey *)r_pubkey, (SigningKey *)d_prvkey, 1);
      int testtype = rand() % 2;
      DecryptWrapper *decwrap = new DecryptWrapper((DecryptionKey *)r_prvkey, (testtype) ? ((CheckingKey *)d_pubkey) : ((CheckingKey *)r_pubkey), 1);

      apr_size_t t_interm_size = 0;
      char *t_interm = encwrap->encryptDataBuffer(pseudorandom, randlen, &t_interm_size);
      
      apr_size_t t_output_size = 0;
      char *t_output = decwrap->decryptDataBuffer(t_interm, t_interm_size, &t_output_size);
      
      apr_size_t new_plain_size = t_output_size;
      char *new_plain_data = (char *)malloc(sizeof(char)*new_plain_size);
      memcpy(new_plain_data, t_output, new_plain_size);
      
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
      } else
      {
         printf("EncDecWrapper Block(issame) test %d passed - data length %d\n", i, randlen);
      }
      free(pseudorandom);
      free(new_plain_data);
      
      // Can't delete memory pools before deleting things in the memory pools.
      free(t_output);
      free(t_interm);
      delete decwrap;
      delete encwrap;
   }
   delete d_pubkey; delete d_prvkey;
   delete r_pubkey; delete r_prvkey;
}

